#
# simple implementation of huffman encoding
#
import sys

class HuffmanNode:
    def __init__(self, sym, prob, child_0 = None, child_1 = None):
        self.sym = sym
        self.prob = prob
        self.child_0 = child_0
        self.child_1 = child_1
        if self.child_0 and self.child_1:
            self.depth = max(self.child_0.depth, self.child_1.depth) + 1
        else:
            self.depth = 1
        
    def __repr__(self):
        return '<HuffmanNode %s %f %s %s>' % \
               (self.sym, self.prob, self.child_0, self.child_1)

    def __cmp__(self, other):
        return cmp(self.prob, other.prob)

class TreeBuilder:
    """symfreq is a list of (symbol string, probability freq)

    Algorithm is the 2-queue method.
    """
    def __init__(self, symfreq):
        self.leaves = [HuffmanNode(i[0], i[1]) for i in symfreq]
        self.leaves.sort()
        self.merged = []

    def build(self):
        while len(self.leaves) > 1 or len(self.merged) > 1:
            if len(self.leaves) > 1:
                leaves_prob = self.leaves[0].prob + self.leaves[1].prob
            else:
                leaves_prob = sys.maxint
            if len(self.merged) > 1:
                merged_prob = self.merged[0].prob + self.merged[1].prob
            else:
                merged_prob = sys.maxint

            if leaves_prob < merged_prob:
                self.merge_leaves()
            elif leaves_prob == merged_prob:
                if (max(self.leaves[0].depth, self.leaves[1].depth) >
                    max(self.merged[0].depth, self.merged[1].depth)):
                    self.merge_merged()
                else:
                    self.merge_leaves()
            else:
                self.merge_merged()

        if len(self.leaves) == 1 and len(self.merged) == 1:
            return HuffmanNode(None, self.leaves[0].prob + self.merged[0].prob,
                               self.leaves[0], self.merged[0])
        else:
            assert len(self.leaves) == 0
            assert len(self.merged) == 1
            return self.merged[0]

    def merge_leaves(self):
        hn = HuffmanNode(None, self.leaves[0].prob + self.leaves[1].prob,
                         self.leaves[0], self.leaves[1])
        self.leaves = self.leaves[2:]
        self.merged.append(hn)

    def merge_merged(self):
        hn = HuffmanNode(None, self.merged[0].prob + self.merged[1].prob,
                         self.merged[0], self.merged[1])
        self.merged = self.merged[2:]
        self.merged.append(hn)

    @staticmethod
    def codes(hufftree):
        c = {}
        def tr(huffnode, code):
            if huffnode.sym:
                c[huffnode.sym] = code
            else:
                tr(huffnode.child_0, code + '0')
                tr(huffnode.child_1, code + '1')
        tr(hufftree, '')
        return c

class Huffman:
    def __init__(self, symfreq):
        self.symfreq = symfreq

    def build_encoding(self):
        pass

if __name__ == '__main__':
    sf = [('aaa', 1),
          ('b', 10),
          ('c', 1)]
    tb = TreeBuilder(sf)
    t = tb.build()
    print TreeBuilder.codes(t)

    sf = [(chr(ord('a') + i), 1) for i in range(26)]
    sf.append(('A', 1000))
    tb = TreeBuilder(sf)
    t = tb.build()
    print TreeBuilder.codes(t)
